from vanna.ZhipuAI.ZhipuAI_Chat import ZhipuAI_Chat
from zhipuai import ZhipuAI
from vanna.openai.openai_chat import OpenAI_Chat
from vanna.chromadb.chromadb_vector import ChromaDB_VectorStore
from config import zhipu_key

clientzhipu = ZhipuAI(api_key=zhipu_key)


class ZhipuVanna(ChromaDB_VectorStore, OpenAI_Chat):
    def __init__(self, config=None):
        ChromaDB_VectorStore.__init__(self, config=config)
        ZhipuAI_Chat.__init__(self, config=config)
        self.client = clientzhipu


vn = ZhipuVanna(config={"model": 'glm-4', "api_key": zhipu_key})
vn.max_tokens = 1000
vn.temperature = 1
vn.train(ddl="""
CREATE TABLE IF NOT EXISTS users (
    id INT PRIMARY KEY COMMENT '用户ID',
    username VARCHAR(50) COMMENT '用户名',
    email VARCHAR(100) COMMENT '电子邮件',
    age INT COMMENT '年龄',
    gender VARCHAR(10) COMMENT '性别（男/女）',
    city VARCHAR(50) COMMENT '城市'
) COMMENT='用户信息表' CHARACTER SET=utf8mb4 COLLATE=utf8mb4_unicode_ci;
""")
vn.train(ddl="CREATE TABLE IF NOT EXISTS consumption_record (   id INT PRIMARY KEY COMMENT '消费记录ID',"
             "user_id COMMENT '用户id' ,item_id INT COMMENT '商品id',amount INT COMMENT '数量' ,consumption INT COMMENT '总消费'"
             ") COMMENT='购买记录' CHARACTER SET=utf8mb4 COLLATE=utf8mb4_unicode_ci;  ")
res = vn.generate_sql('男性用户的总消费是多少')

print('\n\n')

print({'sql': res})
